#!/usr/bin/env python3
"""批量处理 downloads 目录中已有的压缩包，解压并处理图片后清理压缩包。"""

import argparse
from pathlib import Path
from typing import List, Set

from download import (
    setup_download_directory,
    get_archive_type,
    extract_archive_file,
    process_images_in_directory,
)

SUPPORTED_ARCHIVE_TYPES = {"zip", "rar", "7z", "tar"}
ARCHIVE_SUFFIXES: List[str] = [
    ".tar.gz",
    ".tar.bz2",
    ".tar.xz",
    ".tar.lz",
    ".tgz",
    ".tbz2",
    ".zip",
    ".rar",
    ".7z",
    ".tar",
]


def collect_archives(download_path: Path, recursive: bool) -> List[Path]:
    """收集需要处理的压缩包文件。"""
    if recursive:
        candidates = [p for p in download_path.rglob("*") if p.is_file()]
    else:
        candidates = []
        for entry in download_path.iterdir():
            if entry.is_file():
                candidates.append(entry)
            elif entry.is_dir() and entry.name.startswith("process_"):
                candidates.extend(child for child in entry.iterdir() if child.is_file())

    archives: List[Path] = []
    seen: Set[Path] = set()

    for candidate in candidates:
        try:
            archive_type = get_archive_type(str(candidate))
        except Exception:
            continue

        if archive_type in SUPPORTED_ARCHIVE_TYPES:
            resolved = candidate.resolve()
            if resolved not in seen:
                seen.add(resolved)
                archives.append(resolved)

    archives.sort()
    return archives


def strip_archive_suffix(archive_path: Path) -> str:
    """根据压缩包名称推导解压后的目录名称。"""
    lower_name = archive_path.name.lower()
    for suffix in ARCHIVE_SUFFIXES:
        if lower_name.endswith(suffix):
            return archive_path.name[: -len(suffix)]
    return archive_path.stem


def directory_has_files(directory: Path) -> bool:
    """判断目录是否已有内容。"""
    return any(directory.iterdir())


def process_single_archive(archive_path: Path, download_path: Path, watermark_path: Path) -> dict:
    """处理单个压缩包，返回处理状态。"""
    print(f"\n📦 正在处理压缩包: {archive_path.name}")

    extract_dir_name = strip_archive_suffix(archive_path)
    extract_dir = download_path / extract_dir_name
    needs_extraction = not extract_dir.exists() or not directory_has_files(extract_dir)

    extracted = False
    if needs_extraction:
        print(f"🔄 解压目标: {extract_dir}")
        extract_result = extract_archive_file(str(archive_path), str(extract_dir), use_lock=False)
        if not extract_result:
            print("❌ 解压失败，跳过图片处理和删除步骤")
            return {"archive": archive_path, "extracted": False, "processed": False, "deleted": False}
        extracted = True
    else:
        print(f"ℹ️ 目标目录已存在，跳过解压: {extract_dir}")

    processed = process_images_in_directory(str(extract_dir), watermark_path=str(watermark_path))
    if not processed:
        print("⚠️ 图片处理失败或水印缺失，保留压缩包以便人工处理")
        return {"archive": archive_path, "extracted": extracted, "processed": False, "deleted": False}

    try:
        archive_path.unlink()
        deleted = True
        print("🗑️ 已删除压缩包")
    except OSError as err:
        deleted = False
        print(f"⚠️ 删除压缩包失败: {err}")

    return {"archive": archive_path, "extracted": extracted, "processed": True, "deleted": deleted}


def main() -> None:
    parser = argparse.ArgumentParser(description="处理 downloads 目录中的压缩包并完成图片加工")
    parser.add_argument("--downloads", type=str, help="自定义下载目录，默认使用 download 脚本的目录")
    parser.add_argument("--watermark", type=str, default="shuiyin.png", help="水印图片路径")
    parser.add_argument("--recursive", action="store_true", help="递归查找下载目录下的压缩包")
    args = parser.parse_args()

    if args.downloads:
        download_path = Path(args.downloads).expanduser().resolve()
    else:
        download_path = Path(setup_download_directory()).resolve()

    if not download_path.exists():
        print(f"❌ 下载目录不存在: {download_path}")
        return

    watermark_path = Path(args.watermark).expanduser()
    if not watermark_path.is_absolute():
        watermark_path = (Path.cwd() / watermark_path).resolve()

    archives = collect_archives(download_path, recursive=args.recursive)
    if not archives:
        print("✅ 未找到需要处理的压缩包")
        return

    print(f"🔍 共找到 {len(archives)} 个压缩包待处理")

    processed_count = 0
    deleted_count = 0

    for archive in archives:
        result = process_single_archive(archive, download_path, watermark_path)
        if result["processed"]:
            processed_count += 1
        if result["deleted"]:
            deleted_count += 1

    print("\n=== 处理完成 ===")
    print(f"✅ 图片处理成功: {processed_count} 个压缩包")
    print(f"🗑️ 删除压缩包: {deleted_count} 个")
    print(f"📁 下载目录: {download_path}")


if __name__ == "__main__":
    main()
